import os
import pandas as pd
from pandasai import Agent


csv_file_path = 'sales_data.csv'
sales_by_country = pd.read_csv(csv_file_path, sep='\|', names=['country', 'revenue'], engine='python')

sales_by_country['revenue'] = sales_by_country['revenue'].str.strip()


sales_by_country['revenue'] = pd.to_numeric(sales_by_country['revenue'], errors='coerce')


print(sales_by_country.isnull().sum())

# 设置 API 密钥
os.environ["PANDASAI_API_KEY"] = "Please output your own key."

# 创建 Agent 对象并进行查询
try:
    agent = Agent(sales_by_country)
    print("Agent created.")
    response = agent.chat('What is the total sales for the top 3 countries by sales?')
    print("Response from chat:", response)
except Exception as e:
    print("An error occurred:", e)
